#rm(list=ls(all=TRUE))

library(mc2d)

###load data: Y enters first, followed by X variables###

ds <- read.delim("prophecyAJPS.txt",header=TRUE)	# Read into memory a tab delimited file and assign it to a data frame ds
attach(ds)    						# make variables available by name 

dat <- ds[c("P76","P106","NumAzar")]
dat<-as.data.frame(na.omit(dat))


## prob. of 1 for randomizing device or group status indicator ##

p <- 0.2635659

Y <- 1*as.numeric(dat[,1]==2&dat[,2]==2) + 2*as.numeric(dat[,1]==1&dat[,2]==2) + 3*as.numeric(dat[,1]==2&dat[,2]==1) + 4*as.numeric(dat[,1]==1&dat[,2]==1) + 5*as.numeric(dat[,1]==2&dat[,2]==3|dat[,1]==2&dat[,2]==99) + 6*as.numeric(dat[,1]==1&dat[,2]==3|dat[,1]==1&dat[,2]==99)
X <- dat[Y!=0,3]
corrupt.treat <- as.numeric(X==1) 
efficiency.treat <- as.numeric(X==2) 


Y <- Y[Y!=0]

dta<-cbind(Y,corrupt.treat,efficiency.treat)
n<-nrow(dta)


## Set up function to calculate parameter values according to E-M algorithm ##


EM <- function(data,par) {

Y<-data[,1]
X<-as.matrix(cbind(1,data[,2:ncol(data)]))

lamT1 <-par[1]
lamL1 <-par[2]
lamT0 <-par[3]
Beta  <-par[4:length(par)]

mu    <- (1+exp(-X%*%Beta))^-1

n.A <- sum(Y==3|Y==4)
Pr.1.3  <- as.numeric(Y>=3)*0 + as.numeric(Y==1)*((p*lamT0*(1-mu))/(p*lamT0*(1-mu) + (1-p)*lamL1*mu)) + as.numeric(Y==2)*(((1-p)*lamT0*(1-mu))/((1-p)*lamT0*(1-mu) + p*lamL1*mu))
Pr.2.4  <- as.numeric(Y>=3)*0 + as.numeric(Y<3)*(1-Pr.1.3)
Pr.5.6  <- ifelse(Y==3|Y==4,1,0)
Pr.7.9  <- as.numeric(Y<=4)*0 + as.numeric(Y==5)*((p*(1-lamT0)*(1-mu))/(p*(1-lamT0)*(1-mu) + (1-p)*(1-lamT1-lamL1)*mu)) + as.numeric(Y==6)*(((1-p)*(1-lamT0)*(1-mu))/((1-p)*(1-lamT0)*(1-mu) + p*(1-lamT1-lamL1)*mu))
Pr.8.10 <- as.numeric(Y<=4)*0 + as.numeric(Y>4)*(1-Pr.7.9)

Pr.T <- Pr.2.4 + Pr.5.6 + Pr.8.10

lamT1.j <- n.A/(n.A + sum(Pr.2.4) + sum(Pr.8.10) )
lamL1.j <- sum(Pr.2.4)/(n.A + sum(Pr.2.4) + sum(Pr.8.10) )
lamT0.j <- sum(Pr.1.3)/(sum(Pr.1.3)+sum(Pr.7.9) )
  
V <- mu*(1-mu)
X.til<-matrix(NA,nrow=nrow(X),ncol=ncol(X))
for (i in 1:n) {X.til[i,]<-V[i]*X[i,]} 

D <- Pr.T - mu

Beta.j <- Beta + solve(t(X)%*%X.til,tol=.Machine$double.eps^2)%*%t(X)%*%D

est.j <- c(lamT1.j,lamL1.j,lamT0.j,Beta.j)
est.j}

tol <- .Machine$double.eps^.5 # tolerance for stopping algorithm

M<-3000
K<-ncol(dta)
init<-c(.2,.2,.2,rep(0,K)) # initial parameter estimates
pars.old <- EM(dta,init)


for (j in 1:M){

pars.new <- EM(dta,pars.old)
abs.diff <- abs(pars.new-pars.old)
if (max(abs.diff) < tol ) break
pars.old <- pars.new}

iter=j

EM.est<-pars.new

lamT1.mle <-pars.new[1]
lamL1.mle <-pars.new[2]
lamT0.mle <-pars.new[3]
Beta.mle  <-pars.new[4:length(EM.est)]


###Bootstrap Uncertainty Estimates###

n.boot <- 600 					#number of bootstrap samples#
pars.boot <- matrix(NA,nrow=n.boot,ncol=length(pars.new))
delta <-rep(NA,n.boot)
E.mu.1 <-rep(NA,n.boot)
E.mu.0 <-rep(NA,n.boot)

##draw bootstrap samples of X, simulate outcomes using fitted model, then estimate parameter vector for each sample##

for (k in 1:n.boot) {

index<-sample(1:n,size=n,replace=T)
X.boot<-dta[index,2:ncol(dta)]
X.b<-as.matrix(cbind(1,X.boot))

mu.boot <- (1+exp(-X.b%*%Beta.mle))^-1

gam1 <- p*lamT0.mle*(1-mu.boot) + (1-p)*lamL1.mle*mu.boot
gam2 <- (1-p)*lamT0.mle*(1-mu.boot) + p*lamL1.mle*mu.boot
gam3 <- (1-p)*lamT1.mle*mu.boot
gam4 <- p*lamT1.mle*mu.boot
gam5 <- p*(1-lamT0.mle)*(1-mu.boot) + (1-p)*(1-lamT1.mle-lamL1.mle)*mu.boot
gam6 <- (1-p)*(1-lamT0.mle)*(1-mu.boot) + p*(1-lamT1.mle-lamL1.mle)*mu.boot


Mprob <- cbind(gam1,gam2,gam3,gam4,gam5,gam6)

Y.b <- rmultinomial(n=n,size=1,Mprob)
Y.boot <- 1*Y.b[,1] + 2*Y.b[,2] + 3*Y.b[,3] + 4*Y.b[,4] + 5*Y.b[,5] + 6*Y.b[,6]

dat.boot<-cbind(Y.boot,X.boot)

pars.oldB <- EM(dat.boot,init)


for (j in 1:M){

pars.newB <- EM(dat.boot,pars.oldB)
abs.diffB <- abs(pars.newB-pars.oldB)
if (max(abs.diffB) < tol ) break

pars.oldB <- pars.newB}

pars.boot[k,] <- pars.newB

X.1 <- cbind(1,1,0)
X.0 <- cbind(1,0,0)

B.new <- pars.newB[4:length(pars.newB)]
mu.1    <- (1+exp(-X.1%*%B.new))^-1
mu.0    <- (1+exp(-X.0%*%B.new))^-1
E.mu.1[k]<-mean(mu.1)
E.mu.0[k]<-mean(mu.0)
delta[k]<-E.mu.1[k]-E.mu.0[k]

}

V       <- cov(pars.boot)
SEs     <- sqrt(diag(V))
lamT1.CI<- quantile(pars.boot[,1],c(.025,.975))
lamL1.CI<- quantile(pars.boot[,2],c(.025,.975))
lamT0.CI<- quantile(pars.boot[,3],c(.025,.975))

ITT    <- mean(delta)
ITT.se <-sd(delta)
ITT.CI <- quantile(delta,c(.025,.975))

mu1.est <- mean(E.mu.1)
mu1.se  <- sd(E.mu.1)
mu0.est <- mean(E.mu.0)
mu0.se  <- sd(E.mu.0)

Beta.CI <- matrix(NA,nrow=K,ncol=2)
for (z in 1:K){
Beta.CI[z,1:2]<-quantile(pars.boot[,3+z],c(.025,.975))}

conf.int<- rbind(lamT1.CI,lamL1.CI,lamT0.CI,Beta.CI)

out<-cbind(EM.est,SEs,conf.int)
colnames(out)<-c("estimate","SE","2.5%","97.5%")
rownames(out)<-c("lamT1","lamL1","lamT0","constant",colnames(dta[,2:ncol(dta)]))

print(list(parameter.summary=out,iterations=iter,tol=tol,N.obs=n,intent.to.treat=ITT,intent.to.treat.SE=ITT.se,intent.to.treat.CI=ITT.CI))

 
